import torch
import numpy as np
from configs.config import global_var


# 实际图片的标定转为训练时的标签
def real2label(angle, translation, mode):
    if global_var.initial_position == '标准正位' or mode == '标准正位':
        angle_x = angle[0] / 5
        angle_y = (angle[1] - 270) / 5
        angle_z = (angle[2] - 90) / 5
        tx = translation[0] / 25
        ty = translation[1] / 50
        tz = translation[2] / 25
        target = np.array([angle_x, angle_y, angle_z, tx, ty, tz])
        target = torch.from_numpy(target).float().cuda()
        return target


def label2real(angle_noise, translation_noise, batch_size):
    if global_var.initial_position == '标准正位':
        # 标准正位：（0±5， 270±10， 90±5）, t(0±25， 0±50， 0±25)
        pre_alpha = (angle_noise[:, 0] * 5).reshape(batch_size, 1)
        pre_beta = (angle_noise[:, 1] * 10).reshape(batch_size, 1)
        pre_theta = (angle_noise[:, 2] * 5).reshape(batch_size, 1)
        pre_tx = (translation_noise[:, 0] * 25).reshape(batch_size, 1)
        pre_ty = (translation_noise[:, 1] * 50).reshape(batch_size, 1)
        pre_tz = (translation_noise[:, 2] * 25).reshape(batch_size, 1)
    elif global_var.initial_position == '标准侧位':
        # 标准侧位：（0±10,90±5，180±5）, t(0±50， 0±25， 0±25)
        pre_alpha = (angle_noise[:, 0] * 10).reshape(batch_size, 1)
        pre_beta = (angle_noise[:, 1] * 5).reshape(batch_size, 1)
        pre_theta = (angle_noise[:, 2] * 5).reshape(batch_size, 1)
        pre_tx = (translation_noise[:, 0] * 50).reshape(batch_size, 1)
        pre_ty = (translation_noise[:, 1] * 25).reshape(batch_size, 1)
        pre_tz = (translation_noise[:, 2] * 25).reshape(batch_size, 1)
    else:
        raise ValueError("非'标准正位'或'标准侧位'")
    return pre_alpha, pre_beta, pre_theta, pre_tx, pre_ty, pre_tz


def rt2noised_rt(pre_r, mode):
    """将原本的rt改成标准位+noise的形式，方便投影生成"""
    angle_noise = np.array([0, 0, 0]).astype(np.float64)
    if mode == "标准正位":
        angle_noise[0] = pre_r[0]
        angle_noise[1] = pre_r[1] - 270
        angle_noise[2] = pre_r[2] - 90
    elif mode == "标准正位":
        angle_noise[0] = pre_r[0]
        angle_noise[1] = pre_r[1] - 90
        angle_noise[2] = pre_r[2] - 180
    return angle_noise
